import os
import torch
import torchvision
from PIL import Image

from model import *

# 选择设备
device = torch.device("cuda")

# 加载训练好的模型
model = torch.load("best_model.pth")
model = model.to(device)

# 选择验证的图片
root_path = "./imgs"
image_name = "3.png"
image_path = os.path.join(root_path, image_name)

# 读取图片
image = Image.open(image_path).convert("L")

# 用与训练过程相同的transforms进行处理
transforms = torchvision.transforms.Compose([
    torchvision.transforms.Resize((28, 28)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5,), (0.5,))
])
image = transforms(image)

# 进行转换使与模型输入相同
image = torch.reshape(image, (1, 1, 28, 28))
image = image.to(device)

# 进行预测
output = model(image)
label = output.argmax(1).item()

# 输出预测结果
print(label)
